module math_mod

  use const_mod
#ifdef HAS_MKL
  use blas95, only: GEMM
  use lapack95, only: SYEV
#endif

  private
  
  public cross_product
  public matrix_inverse
  public matrix_mutiply
  public symmetric_eigen_solver

contains

  function cross_product(x, y) result(res)

    real(16), intent(in) :: x(3)
    real(16), intent(in) :: y(3)
    real(16) res(3)

    res(1) = x(2) * y(3) - x(3) * y(2)
    res(2) = x(3) * y(1) - x(1) * y(3)
    res(3) = x(1) * y(2) - x(2) * y(1)

  end function cross_product
    
  subroutine matrix_inverse(n, A, B)

    integer , intent(in   ) :: n
    real(r8), intent(inout) :: A(n,n)
    real(r8), intent(  out) :: B(n,n)

    real(r8) C(n,n)
    integer ipiv(n)
    integer i, j, k

    C = 0.0d0
    do i = 1, n
      C(i,i) = 1.0d0
    end do

    call gaussian_elimination(n, A, ipiv)

    do i = 1, n - 1
      do j = i + 1, n
        do k = 1, n
          C(ipiv(j),k) = C(ipiv(j),k) - A(ipiv(j),i) * C(ipiv(i),k)
        end do
      end do
    end do

    do i = 1, n
      B(n,i) = C(ipiv(n),i) / A(ipiv(n),n)
      do j = n - 1, 1, -1
        B(j,i) = C(ipiv(j),i)
        do k = j + 1, n
          B(j,i) = B(j,i) - A(ipiv(j),k) * B(k,i)
        end do
        B(j,i) = B(j,i) / A(ipiv(j),j)
      end do
    end do

  end subroutine matrix_inverse

  subroutine gaussian_elimination(n, A, ipiv)

    integer , intent(in   ) :: n
    real(r8), intent(inout) :: A(n,n)
    integer , intent(  out) :: ipiv(n)

    real(r8) c1, c(n), pi, pi1, pj
    integer i, j, k, itmp

    do i = 1, n
      ipiv(i) = i
    end do

    do i = 1, n
      c1 = 0.0d0
      do j = 1, n
        c1 = max(c1, abs(A(i,j)))
      end do
      c(i) = c1
    end do

    do j = 1, n - 1
      pi1 = 0.0
      do i = j, n
        pi = abs(A(ipiv(i),j)) / c(ipiv(i))
        if (pi > pi1) then
          pi1 = pi
          k   = i
        end if
      end do

      itmp    = ipiv(j)
      ipiv(j) = ipiv(k)
      ipiv(k) = itmp
      do i = j + 1, n
        pj = A(ipiv(i),j) / A(ipiv(j),j)
        A(ipiv(i),j) = pj
        do k = j + 1, n
          A(ipiv(i),k) = A(ipiv(i),k) - pj * A(ipiv(j),k)
        end do
      end do
    end do

  end subroutine gaussian_elimination

  subroutine matrix_mutiply(transA, transB, A, B, C)

    character(1), intent(in) :: transA
    character(1), intent(in) :: transB
    real(8), intent(in) :: A(:,:)
    real(8), intent(in) :: B(:,:)
    real(8), intent(out) :: C(:,:)

    integer m, n, k, lda, ldb, ldc

#ifdef HAS_MKL
    call GEMM(A, B, C, transA, transB)
#else
    m = merge(size(A, 1), size(A, 2), transA == 'N')
    n = merge(size(B, 2), size(B, 1), transB == 'N')
    k = merge(size(A, 2), size(A, 1), transA == 'N')
    lda = merge(max(1, m), max(1, k), transA == 'N')
    ldb = merge(max(1, k), max(1, n), transB == 'N')
    ldc = max(1, m)
    call DGEMM(transA, transB, m, n, k, 1.0d0, A, lda, B, ldb, 0.0d0, C, ldc)
#endif

  end subroutine matrix_mutiply

  subroutine symmetric_eigen_solver(A, eigvals, eigvecs)

    real(8), intent(in) :: A(:,:)
    real(8), intent(inout) :: eigvals(:)
    real(8), intent(inout), optional :: eigvecs(:,:)

    real(8), allocatable :: B(:,:)
    real(8), allocatable :: W(:)
    real(8) tmp(1)
    integer n, lda, ierr

#ifdef HAS_MKL
    if (present(eigvecs)) then
      eigvecs = A
      call SYEV(eigvecs, eigvals, 'V', 'U')
    else
      allocate(B(size(A,1),size(A,2)))
      B = A
      call SYEV(B, eigvals, 'V', 'U')
      deallocate(B)
    end if
#else
    if (present(eigvecs)) then
      lda = size(A, 1)
      n   = size(A, 2)
      eigvecs = A
      call DSYEV('V', 'U', n, eigvecs, lda, eigvals, tmp, -1, ierr)
      allocate(W(int(tmp(1))))
      call DSYEV('V', 'U', n, eigvecs, lda, eigvals, W, size(W), ierr)
      deallocate(W)
    else
      allocate(B(size(A,1),size(A,2)))
      B = A
      lda = size(A, 1)
      n   = size(A, 2)
      call DSYEV('V', 'U', n, B, lda, eigvals, tmp, -1, ierr)
      allocate(W(int(tmp(1))))
      call DSYEV('V', 'U', n, B, lda, eigvals, W, size(W), ierr)
      deallocate(B)
    end if
#endif

  end subroutine symmetric_eigen_solver

end module math_mod
